-- uncomment the module() line when you have 
-- written and debugged the code in this 
-- file.  It allows *another* file to get 
-- these functions using 
-- pq = require("priority_queues") 
-- 
-- and fn names like pq.build_heap(a)

module(..., package.seeall) 

function parent (i) 
   if i>1
   then return math.floor(i/2) 
   else return nil
   end
end
   
function left_child (i)
   return i*2
end 

function right_child (i)
   return (i*2)+1 
end 

-- If h is a heap, then h.pref(v,w)
-- means that v is preferred for a position 
-- above w in the heap.  

-- h.loc(v,i) should be called whenever 
-- value v is being placed into position i
-- in heap h.  

function make_empty_heap(above, locate) 
   return { pref = above, 
	    loc = locate }
end



--[[

heapify_up(h,i) takes an almost-heap 
h, that is, an array in which for any 
other index j where j not= i, 
h[j] belongs below its parent. 

If h[i] belongs *above* its parent, 
then heapify_up does this interchange, 
and then continues to
heapify_up(h,parent(i)).  

--]]

function heapify_up(h,i)
	
	local p = parent(i)
	
	--print("heapify_up", i, p, h[i],  h[p] )
	--print( unpack(h))
	
	if p then
		if h.pref(h[i], h[p]) then
			h[i], h[p] = h[p], h[i]
			heapify_up(h,p)
		end	
	end
	
	--if p == nil or h[i] == nil or not(h.pref(h[i],h[p])) then
		
	--end

end


--[[

heapify_down takes a broken heap and 
takes one step to repair it.  In particular, 
its arguments are an array a and an index 
i, such that: 

1.  the left and right subtrees dominated by i 
    are already heaps; 

2.  a[i] belongs below its parent if any.  

If a[i] belongs below than either of its children, 
heapify_down flips it with whichever child would
belong above the other.  

It recursively repairs any damage this may have 
caused to the heap property of this child.

--]]

function heapify_down (h,i) 

	local left = left_child (i)
	local right = right_child (i)
	
	--print("Down: ", i, left, right, h[i], h[left], h[right] )
	--print( unpack( h ) )
	
	if h[left] ~=nil and ( h[i] == nil or h.pref(h[left], h[i]) ) then
		h[i], h[left] = h[left], h[i]
		heapify_down(h,left)
	
	end
	
	if h[right] ~=nil and ( h[i] == nil or h.pref(h[right], h[i]) ) then
		h[i], h[right] = h[right], h[i]
		heapify_down(h,right)
	
	end	
	
end  




--[[

build_heap takes any array of numbers and 
optionally functions for comparison and 
recording updates.  

build_heap initializes a heap in which each
parent belongs above its children.

The local defns of above and locate supply
a default value to make_empty_heap 
if build_heap is called without supplying
a value for g, or without supplying
values for f,g.  

The default for above calls >, and 
the default for locate does nothing.

--]]

function build_heap (a, f, g) 
	local above = f or
		function (v,u)
			return v>u
		end 
	local locate = g or
		function (v,u)
		end
	local h =  make_empty_heap(above, locate) 
   
	for k, v in ipairs(a) do
		insert( h, v )
	end
	
	for i = math.floor(#a/2), 1, -1 do
		heapify_down (h, i) 
	end
	
   return h
end


function insert(h,v) 
   local i = #h+1 
   --   while h[i] do i=i+1 end 
   h[i] = v 
   heapify_up(h,i)
end 


function find_top(h) 
   return h[1]
end 

function delete(h,i) 
   h[i] = nil 
   heapify_down(h,i) 
end 

function extract_top(h) 
   local v = find_top(h) 
   delete(h,1) 
   return v 
end 

function check_heap (h,i) 
   for j = 1,i do 
      local l = left_child(j) 
      local r = right_child(j)
      if (not(h[l]) or not(h.pref(h[l], h[j])))
      and (not(h[r]) or not(h.pref(h[r],h[j])))
      then 			
	 -- No problem here! 
      else
	 error(string.format("check_heap failed at %d\n", j))
      end 
   end 
end 